import matplotlib.pyplot as plt
import matplotlib.patches as patches


def create_architecture_diagram():
    # 创建画布
    fig, ax = plt.subplots(figsize=(12, 8))

    # 定义模块位置和大小
    modules = {
        'Backbone': (0.1, 0.5, 0.15, 0.2),
        'LAM-FPN': (0.35, 0.5, 0.15, 0.2),
        'ADEM': (0.35, 0.2, 0.15, 0.15),
        'LCM': (0.6, 0.5, 0.15, 0.2),
        'Detection Head': (0.85, 0.5, 0.15, 0.2)
    }

    # 定义模块颜色
    colors = {
        'Backbone': '#AED6F1',
        'LAM-FPN': '#F5B7B1',
        'ADEM': '#A2D9CE',
        'LCM': '#D7BDE2',
        'Detection Head': '#FAD7A0'
    }

    # 绘制模块框
    for name, (x, y, w, h) in modules.items():
        rect = patches.Rectangle(
            (x, y), w, h,
            facecolor=colors[name],
            edgecolor='black',
            linewidth=1,
            alpha=0.7
        )
        ax.add_patch(rect)
        ax.text(x + w / 2, y + h / 2, name, ha='center', va='center')

    # 添加连接箭头
    arrows = [
        ((0.25, 0.6), (0.35, 0.6), 'Features'),
        ((0.5, 0.6), (0.6, 0.6), 'Enhanced\nFeatures'),
        ((0.75, 0.6), (0.85, 0.6), 'Final\nFeatures'),
        ((0.425, 0.5), (0.425, 0.35), 'Feature\nSharing'),
        ((0.425, 0.2), (0.85, 0.5), 'Density\nGuidance')
    ]

    for (start_x, start_y), (end_x, end_y), label in arrows:
        ax.annotate(
            '',
            xy=(end_x, end_y),
            xytext=(start_x, start_y),
            arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0.1')
        )
        # 计算标签位置
        mid_x = (start_x + end_x) / 2
        mid_y = (start_y + end_y) / 2
        ax.text(mid_x, mid_y, label, ha='center', va='center')

    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis('off')
    plt.title('Enhanced DiffusionDet Architecture', pad=20)

    # 保存图片
    plt.savefig('density_aware_mechanism_architecture.pdf', bbox_inches='tight', dpi=300)
    plt.close()


create_architecture_diagram()
